import streamlit as st
import cv2
import numpy as np
import os
import sqlite3
from datetime import datetime
import time


# 初始化数据库 - 修复版本
def init_db():
    conn = sqlite3.connect("attendance.db")
    cursor = conn.cursor()

    # 先检查表是否存在，如果存在则删除（确保结构正确）
    cursor.execute("DROP TABLE IF EXISTS attendance")

    # 创建正确的表结构
    cursor.execute("""
        CREATE TABLE attendance (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name VARCHAR(20) NOT NULL,
            date DATE NOT NULL,
            clocking_time TIME NOT NULL
        )
    """)
    conn.commit()
    conn.close()


# 录入考勤记录
def record_attendance(name):
    conn = sqlite3.connect("attendance.db")
    cursor = conn.cursor()
    current_date = datetime.now().strftime("%Y-%m-%d")
    current_time = datetime.now().strftime("%H:%M:%S")

    try:
        cursor.execute(
            "INSERT INTO attendance (name, date, clocking_time) VALUES (?, ?, ?)",
            (name, current_date, current_time)
        )
        conn.commit()
        st.success(f"✅ 已记录 {name} 的考勤 - {current_time}")
    except sqlite3.Error as e:
        st.error(f"数据库错误: {str(e)}")
    finally:
        conn.close()


# 获取考勤记录
def get_attendance_records():
    conn = sqlite3.connect("attendance.db")
    cursor = conn.cursor()
    try:
        cursor.execute("SELECT name, date, clocking_time FROM attendance ORDER BY date DESC, clocking_time DESC")
        records = cursor.fetchall()
        return records
    except sqlite3.Error as e:
        st.error(f"查询错误: {str(e)}")
        return []
    finally:
        conn.close()


# 改进的人脸识别函数
def recognize_face(frame, known_faces_dir="./faces"):
    # 转换为灰度图像
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 加载人脸检测器
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
    faces = face_cascade.detectMultiScale(gray, 1.1, 5, minSize=(50, 50))

    recognized_name = "未知"

    if len(faces) > 0:
        # 只处理第一个人脸
        x, y, w, h = faces[0]
        face_roi = gray[y:y + h, x:x + w]

        # 与已知人脸进行比较
        best_match = None
        best_score = 0

        # 检查已知人脸目录是否存在
        if os.path.exists(known_faces_dir) and os.listdir(known_faces_dir):
            for file in os.listdir(known_faces_dir):
                if file.endswith(('.png', '.jpg', '.jpeg')):
                    known_face_path = os.path.join(known_faces_dir, file)
                    known_face = cv2.imread(known_face_path, cv2.IMREAD_GRAYSCALE)

                    if known_face is not None:
                        # 调整尺寸以匹配
                        known_face = cv2.resize(known_face, (w, h))

                        # 使用模板匹配
                        result = cv2.matchTemplate(face_roi, known_face, cv2.TM_CCOEFF_NORMED)
                        _, max_val, _, _ = cv2.minMaxLoc(result)

                        if max_val > best_score and max_val > 0.65:  # 降低匹配阈值
                            best_score = max_val
                            best_match = os.path.splitext(file)[0]

        if best_match:
            recognized_name = best_match
            # 在图像上绘制矩形和名称
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            cv2.putText(frame, f"{recognized_name} ({best_score:.2f})", (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        else:
            # 绘制检测到但未识别的人脸
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            cv2.putText(frame, "未知", (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    else:
        # 没有检测到人脸
        cv2.putText(frame, "未检测到人脸", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    return frame, recognized_name


# 主应用程序
def main():
    st.title("人脸识别考勤系统")

    # 初始化数据库
    init_db()

    # 创建选项卡
    tab1, tab2, tab3 = st.tabs(["人脸录入", "考勤识别", "考勤记录"])

    with tab1:
        st.header("人脸信息录入")

        # 创建保存人脸的目录
        if not os.path.exists("./faces"):
            os.makedirs("./faces")

        # 使用摄像头输入
        img_file_buffer = st.camera_input("请面对摄像头拍照")

        if img_file_buffer is not None:
            # 将图像转换为OpenCV格式
            bytes_data = img_file_buffer.getvalue()
            cv2_img = cv2.imdecode(np.frombuffer(bytes_data, np.uint8), cv2.IMREAD_COLOR)

            if cv2_img is not None:
                # 转换为灰度图像进行人脸检测
                gray = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2GRAY)
                face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
                faces = face_cascade.detectMultiScale(gray, 1.1, 5, minSize=(50, 50))

                if len(faces) > 0:
                    x, y, w, h = faces[0]
                    face_roi = cv2_img[y:y + h, x:x + w]

                    # 显示检测到的人脸
                    col1, col2 = st.columns(2)
                    with col1:
                        st.image(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB), caption="原始图像")
                    with col2:
                        st.image(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB), caption="检测到的人脸", width=200)

                    # 输入姓名
                    name = st.text_input("请输入姓名")

                    # 显示已录入的人脸列表
                    if os.path.exists("./faces") and os.listdir("./faces"):
                        st.info("已录入的人脸: " + ", ".join([os.path.splitext(f)[0] for f in os.listdir("./faces") if
                                                              f.endswith(('.png', '.jpg', '.jpeg'))]))

                    if st.button("保存人脸") and name:
                        # 保存人脸图像
                        face_path = os.path.join("./faces", f"{name}.png")
                        cv2.imwrite(face_path, face_roi)
                        st.success(f"✅ 已保存 {name} 的人脸信息")
                        st.rerun()
                else:
                    st.warning("❌ 未检测到人脸，请重新拍照")
            else:
                st.error("无法处理图像")

    with tab2:
        st.header("人脸考勤识别")

        # 开始/停止识别按钮
        if 'recognizing' not in st.session_state:
            st.session_state.recognizing = False
        if 'last_recognition' not in st.session_state:
            st.session_state.last_recognition = None
        if 'last_recognition_time' not in st.session_state:
            st.session_state.last_recognition_time = 0

        col1, col2 = st.columns(2)
        with col1:
            if st.button("开始识别", type="primary", disabled=st.session_state.recognizing):
                st.session_state.recognizing = True
                st.session_state.last_recognition = None
                st.session_state.last_recognition_time = 0

        with col2:
            if st.button("停止识别", disabled=not st.session_state.recognizing):
                st.session_state.recognizing = False

        if st.session_state.recognizing:
            st.info("🔴 识别中... 请面对摄像头")

            # 实时摄像头输入
            camera = st.camera_input("实时摄像头", key="recognition_camera")

            if camera is not None:
                # 将图像转换为OpenCV格式
                bytes_data = camera.getvalue()
                frame = cv2.imdecode(np.frombuffer(bytes_data, np.uint8), cv2.IMREAD_COLOR)

                if frame is not None:
                    # 进行人脸识别
                    processed_frame, name = recognize_face(frame)

                    # 显示处理后的图像
                    st.image(cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB), caption="识别结果")

                    # 如果识别到已知人脸且距离上次记录超过5秒
                    current_time = time.time()
                    if (name != "未知" and
                            (st.session_state.last_recognition != name or
                             current_time - st.session_state.last_recognition_time > 5)):
                        # 记录考勤
                        record_attendance(name)

                        # 更新最后识别信息
                        st.session_state.last_recognition = name
                        st.session_state.last_recognition_time = current_time
                else:
                    st.error("无法处理摄像头图像")
        else:
            st.info("请点击「开始识别」按钮")

    with tab3:
        st.header("考勤记录")

        # 显示考勤记录
        records = get_attendance_records()

        if records:
            # 显示统计信息
            today = datetime.now().strftime("%Y-%m-%d")
            today_count = len([r for r in records if r[1] == today])
            total_count = len(records)

            col1, col2 = st.columns(2)
            col1.metric("今日考勤", today_count)
            col2.metric("总记录数", total_count)

            # 转换为DataFrame格式显示
            import pandas as pd
            df = pd.DataFrame(records, columns=["姓名", "日期", "时间"])
            st.dataframe(df)

            # 提供下载选项
            csv = df.to_csv(index=False, encoding='utf-8-sig')
            st.download_button(
                label="下载考勤记录",
                data=csv,
                file_name="考勤记录.csv",
                mime="text/csv"
            )
        else:
            st.info("暂无考勤记录")


if __name__ == "__main__":
    main()